#!/bin/env python3
import sys
import os

def update(old_sif,new_sif,post_cmd):

    update_def=f"""Bootstrap:localimage
From: {old_sif}
%setup
%files
%post
    {post_cmd}
%environment
%runscript
    """
    with open("update.def","w") as f:
        f.writelines(update_def)
    os.system(f"cd /CLQCD/containers; singularity build --fakeroot {new_sif} /opt/containers/update.def")

def main():
    post_lines_update_source_list="""
    # 备份原有源
    cp /etc/apt/sources.list /etc/apt/sources.list.bak

    # 替换为阿里云镜像源
    sed -i 's|http://mirrors.tuna.tsinghua.edu.cn|https://mirrors.aliyun.com|g' /etc/apt/sources.list

    # 更新软件包索引
    apt-get update -y

    """
    old_sif = sys.argv[1]
    new_sif = sys.argv[2]
    #post_cmd = sys.argv[3]

    # pyquda
    post_cmd="pip3 install gmpy2;pip3 install -U setuptools;pip3 install -U pip;cd /src/PyQuda;python3 -m pip install ."
    # sources.list
    post_cmd=post_lines_update_source_list
    update(old_sif,new_sif,post_cmd)

if __name__== "__main__":
    main()